//
// Created by lidongwei on 2023/6/14.
//

#include "pretreatment.h"


////////////////////////////////////仿射变换////////////////////////////////////////
__device__ void affine_project(float *matrix, int x, int y, float *proj_x, float *proj_y) {

    // matrix
    // m0, m1, m2
    // m3, m4, m5
    *proj_x = matrix[0] * x + matrix[1] * y + matrix[2];
    *proj_y = matrix[3] * x + matrix[4] * y + matrix[5];
}

__global__ void kernel_Affine_transformation(uint8_t *in_device, uint8_t *out_device, float *matrix,
                                             cv::Size output_size, cv::Size img_size, int mode) {//仿射变换的核函数
    int dx = blockDim.x * blockIdx.x + threadIdx.x;
    int dy = blockDim.y * blockIdx.y + threadIdx.y;
    //printf("%d,%d,%d,%d,%d,%d\n",blockDim.x,blockIdx.x,threadIdx.x,blockDim.y,blockIdx.y,threadIdx.y);
    if (dx >= output_size.width || dy >= output_size.height) {
        printf("%d,%d,out", dx, dy);
        return;
    }

    int c0 = 114, c1 = 114, c2 = 114;//灰度值填充，BGR
    //printf("%d,%d,%d",c0,c1,c2);
    float src_x = 0, src_y = 0;
    affine_project(matrix, dx, dy, &src_x, &src_y);//在原图找到位置
    //双线性插值
    if (src_x < -1 || src_x >= img_size.width || src_y < -1 || src_y >= img_size.height) {
        return;
        // out of range
        // src_x < -1时，其高位high_x < 0，超出范围
        // src_x >= -1时，其高位high_x >= 0，存在取值
    } else {
        int x_low = floorf(src_x);
        int y_low = floorf(src_y);
        int x_high = x_low + 1;
        int y_high = y_low + 1;

        float lx = src_x - x_low;
        float ty = src_y - y_low;
        float rx = x_high - src_x;
        float dy = y_high - src_y;

        float w1 = lx * ty, w2 = rx * ty, w3 = lx * dy, w4 = rx * dy;
        uint8_t const_values[] = {114, 114, 114};
        uint8_t *v1 = const_values;
        uint8_t *v2 = const_values;
        uint8_t *v3 = const_values;
        uint8_t *v4 = const_values;

        if (y_low >= 0) {
            if (x_low >= 0)
                v1 = in_device + y_low * img_size.width * 3 + x_low * 3;

            if (x_high < img_size.width)
                v2 = in_device + y_low * img_size.width * 3 + x_high * 3;
        }

        if (y_high < img_size.height) {
            if (x_low >= 0)
                v3 = in_device + y_high * img_size.width * 3 + x_low * 3;

            if (x_high < img_size.width)
                v4 = in_device + y_high * img_size.width * 3 + x_high * 3;
        }

        c0 = floorf((v1[0] * w4 + v2[0] * w3 + v3[0] * w2 + v4[0] * w1 + 0.5f));//使用地址机制进行取数
        c1 = floorf((v1[1] * w4 + v2[1] * w3 + v3[1] * w2 + v4[1] * w1 + 0.5f));
        c2 = floorf((v1[2] * w4 + v2[2] * w3 + v3[2] * w2 + v4[2] * w1 + 0.5f));

    }
    /////////////////RGB--BGR/////////////////////////
    uint8_t *out = out_device + dy * output_size.width * 3 + dx * 3;
    if (mode == 1) out[0] = c0, out[1] = c1, out[2] = c2;//BGR
    else if (mode == 2) out[0] = c2, out[1] = c1, out[2] = c0;//RGB

}

cv::Mat handle::pretreatment::Affine_transformation(cv::Mat img, cv::Size output_size, float theta, int mode = 1) {
//    printf("%d", mode);
    uint8_t *in_device = nullptr;
    uint8_t *out_device = nullptr;
    theta = theta * M_PI / 180.0;

    size_t in_size = img.cols * img.rows * 3;
    size_t out_size = output_size.width * output_size.height * 3;
    cv::Size img_size = cv::Size(img.cols, img.rows);
    cv::Mat out_img(output_size, CV_8UC3);

    checkRuntime(cudaMalloc(&in_device, in_size));//在GPU申请两块内存，分别存放转换前和转换后的图片
    checkRuntime(cudaMalloc(&out_device, out_size));
    checkRuntime(cudaMemcpy(in_device, img.data, in_size, cudaMemcpyHostToDevice));

    //定义仿射矩阵
    //    [
    //    scale*cos(theta),   -sin(theta)*scale,    -scale * from.width * 0.5 + to.width * 0.5
    //    sin(theta)*scale,   scale*cos(theta) ,    -scale * from.height * 0.5 + to.height * 0.5
    //           0        ,           0        ,                     1
    //    ]
    float scale = min(output_size.width / (float) img.cols, output_size.height / (float) img.rows);
    float sin_ = sin(theta);
    float cos_ = cos(theta);
    float ox = 0, oy = 0;
    ox = -(scale * img.cols) / 2 + output_size.width / 2;
    oy = -(scale * img.rows) / 2 + output_size.height / 2;
    cv::Mat S_warp_matrix = (cv::Mat_<float>(3, 3) << scale, 0, 0, 0,
            scale, 0, 0, 0, 1);//缩放矩阵
    cv::Mat T_warp_matrix = (cv::Mat_<float>(3, 3) << 1, 0, ox, 0,
            1, oy, 0, 0, 1);//平移矩阵
    cv::Mat R_warp_matrix = (cv::Mat_<float>(3, 3) << 1, 0, img.cols / 2, 0, 1, img.rows / 2, 0, 0, 1) *
                            (cv::Mat_<float>(3, 3) << cos_, -sin_, 0, sin_, cos_, 0, 0, 0, 1) *
                            (cv::Mat_<float>(3, 3) << 1, 0, -img.cols / 2, 0, 1, -img.rows / 2, 0, 0, 1);//旋转矩阵

    cv::Mat warp_matrix = T_warp_matrix * S_warp_matrix * R_warp_matrix;//旋转->缩放->平移
    //求仿射矩阵逆矩阵
    cv::Mat warp_inverse_matrix_translation = cv::Mat_<float>(3, 3);
    cv::invert(warp_matrix, warp_inverse_matrix_translation);

    float *matrix = nullptr;//将仿射矩阵逆矩阵放到GPU上，不然无法进行访问调用
    checkRuntime(cudaMalloc(&matrix, 9 * sizeof(float)));
    checkRuntime(cudaMemcpy(matrix, warp_inverse_matrix_translation.data, 9 * sizeof(float), cudaMemcpyHostToDevice));

    //初始化cuda核
    dim3 block_size(32, 32);
    dim3 grid_size((output_size.width + 31) / 32, (output_size.height + 31) / 32);//31是为了向上取整

    //调用核函数
    kernel_Affine_transformation<<<grid_size, block_size, 0, nullptr>>>(
            in_device, out_device, matrix, output_size, img_size, mode
    );
    //检查核函数运行是否存在错误
    cudaDeviceSynchronize(); // 在核函数执行之后
    cudaError_t error = cudaPeekAtLastError();
    if (error != cudaSuccess) {
        printf("CUDA error: %s\n", cudaGetErrorString(error));
    }

    checkRuntime(cudaMemcpy(out_img.data, out_device, out_size * sizeof(uint8_t), cudaMemcpyDeviceToHost));
    checkRuntime(cudaFree(out_device));
    checkRuntime(cudaFree(in_device));
    checkRuntime(cudaFree(matrix));

    return out_img;
}


/////////////////////////////////////HWC2CHW//////////////////////////////////////////////////
__global__ void kernel_HWC2CHW(uint8_t *data, uint8_t *out_data, int width, int height) {
    int dx = blockDim.x * blockIdx.x + threadIdx.x;
    int dy = blockDim.y * blockIdx.y + threadIdx.y;

    out_data[dy * width + dx + width * height * 0] = data[dy * width * 3 + dx * 3 + 0];
    out_data[dy * width + dx + width * height * 1] = data[dy * width * 3 + dx * 3 + 1];
    out_data[dy * width + dx + width * height * 2] = data[dy * width * 3 + dx * 3 + 2];

}

cv::Mat handle::pretreatment::HWC2CHW(cv::Mat img) {
    int width = img.cols;
    int height = img.rows;
    int channel = img.channels();

    cv::Mat out_data_host(height, width, CV_8UC3);

    uint8_t *data_device = nullptr;
    uint8_t *out_data_device = nullptr;
    checkRuntime(cudaMalloc(&data_device, width * height * channel));
    checkRuntime(cudaMalloc(&out_data_device, width * height * channel));
    checkRuntime(cudaMemcpy(data_device, img.data, width * height * channel, cudaMemcpyHostToDevice));

    //初始化cuda核
    dim3 block_size(32, 32);
    dim3 grid_size((width + 31) / 32, (height + 31) / 32);//31是为了向上取整

    kernel_HWC2CHW<<<grid_size, block_size, 0, nullptr>>>(data_device, out_data_device, width, height);

    checkRuntime(cudaMemcpy(out_data_host.data, out_data_device, width * height * channel, cudaMemcpyDeviceToHost));

    checkRuntime(cudaFree(data_device));
    checkRuntime(cudaFree(out_data_device));

    return out_data_host;
}

/////////////////////////////////////CHW2HWC//////////////////////////////////////////////////
__global__ void kernel_CHW2WHC(uint8_t *data, uint8_t *out_data, int width, int height) {
    int dx = blockDim.x * blockIdx.x + threadIdx.x;
    int dy = blockDim.y * blockIdx.y + threadIdx.y;

    out_data[dy * width * 3 + dx * 3 + 0] = data[dy * width + dx + width * height * 0];
    out_data[dy * width * 3 + dx * 3 + 1] = data[dy * width + dx + width * height * 1];
    out_data[dy * width * 3 + dx * 3 + 2] = data[dy * width + dx + width * height * 2];

}

cv::Mat handle::pretreatment::CHW2HWC(cv::Mat img) {
    int width = img.cols;
    int height = img.rows;
    int channel = img.channels();

    cv::Mat out_data_host(height, width, CV_8UC3);

    uint8_t *data_device = nullptr;
    uint8_t *out_data_device = nullptr;
    checkRuntime(cudaMalloc(&data_device, width * height * channel));
    checkRuntime(cudaMalloc(&out_data_device, width * height * channel));
    checkRuntime(cudaMemcpy(data_device, img.data, width * height * channel, cudaMemcpyHostToDevice));

    //初始化cuda核
    dim3 block_size(32, 32);
    dim3 grid_size((width + 31) / 32, (height + 31) / 32);//31是为了向上取整

    kernel_CHW2WHC<<<grid_size, block_size, 0, nullptr>>>(data_device, out_data_device, width, height);

    checkRuntime(cudaMemcpy(out_data_host.data, out_data_device, width * height * channel, cudaMemcpyDeviceToHost));

    checkRuntime(cudaFree(data_device));
    checkRuntime(cudaFree(out_data_device));

    return out_data_host;
}

////////////////////////////归一化///////////////////////////////
__global__ void
kernel_Normalization(uint8_t *data, float *out_data, int width, int height, int Normalization, int model_channels) {
    int dx = blockDim.x * blockIdx.x + threadIdx.x;
    int dy = blockDim.y * blockIdx.y + threadIdx.y;

    for (int i = 0; i < model_channels; i++) {
        out_data[dy * width + dx + i * width * height] =
                (float) data[dy * width + dx + +i * width * height] / Normalization;
    }

}

std::vector<float> handle::pretreatment::Normalization(cv::Mat img, int Normalization, int model_channels = 3) {
//    auto start=clock();
    int width = img.cols;
    int height = img.rows;
    int channel = img.channels();
    cudaStream_t stream= nullptr;
    cudaStreamCreate(&stream);

    cv::Mat img_chw = handle::pretreatment::HWC2CHW(img);

    std::vector<float> out_data_host(height * width * model_channels);

    uint8_t *data_device = nullptr;
    float *out_data_device = nullptr;
    checkRuntime(cudaMallocAsync(&data_device, width * height * model_channels,stream));
    checkRuntime(cudaMallocAsync(&out_data_device, width * height * model_channels * sizeof(float),stream));
    checkRuntime(cudaMemcpyAsync(data_device, img_chw.data, width * height * channel, cudaMemcpyHostToDevice,stream));

    if (model_channels > 3) {//判断是否有第四通道
        cv::Mat out_data_host_four_channel(height, width, CV_8UC3);
        cv::convertScaleAbs(img, out_data_host_four_channel, 1.5, 20);
        cv::cvtColor(out_data_host_four_channel, out_data_host_four_channel, cv::COLOR_RGB2GRAY);
        checkRuntime(
                cudaMemcpyAsync(data_device + width * height * channel, out_data_host_four_channel.data, width * height * 1,
                           cudaMemcpyHostToDevice,stream));
    }

    //初始化cuda核
    dim3 block_size(32, 32);
    dim3 grid_size((width + 31) / 32, (height + 31) / 32);//31是为了向上取整

    kernel_Normalization<<<grid_size, block_size, 0, stream>>>(data_device, out_data_device, width, height,
                                                                Normalization, model_channels);
    checkRuntime(cudaMemcpyAsync(out_data_host.data(), out_data_device, width * height * model_channels * sizeof(float),
                            cudaMemcpyDeviceToHost,stream));
    cudaStreamSynchronize(stream);

//    auto end=clock();
//    double diff_time=(double)(end-start)/CLOCKS_PER_SEC*1000;
//    printf("整体耗时：%.1f ms\n",diff_time);
    checkRuntime(cudaFree(data_device));
    checkRuntime(cudaFree(out_data_device));

    return out_data_host;
}
